from mindspore import ops, Tensor
from mindspore import dtype as mstype
from mindspore import numpy as mnp


"""
    在处理数据时，我们通常会遇到序列长度不一致的问题。为了统一长度，会使用 <pad> 占位符补齐了一些稍短的文本。
    如："Hello world!" --> <bos> hello world ! <eos> <pad> <pad>
    这些 <pad> 占位符没有任何意义，不应该参与注意力分数计算中。为此我们在注意力中加入了 padding 掩码，
    即识别输入序列中的 <pad> 占位符，保证计算时这些位置对应的注意力分数为0。
"""


def get_attn_pad_mask(seq_q: Tensor, seq_k: Tensor, pad_idx: int):
    """注意力掩码：识别序列中的<pad>占位符
    Args:
        seq_q : query序列，shape = [batch size, query_len]
        seq_k : key序列，shape = [batch size, key_len]
        pad_idx : key序列<pad>占位符对应的数字索引
    """
    batch_size, len_q = seq_q.shape
    batch_size, len_k = seq_k.shape

    # 检查 seq_k 中的每个元素，元素值等于 pad_idx的对应位置为True，否则为False；得到一个布尔型的掩码 pad_attn_mask： [batch_size, len_k]
    pad_attn_mask = ops.equal(seq_k, pad_idx)
    # 将 pad_attn_mask 扩展为 [batch_size, 1, len_k] 形状，用于与 seq_q 进行广播
    pad_attn_mask = pad_attn_mask.expand_dims(1)
    # 将 pad_attn_mask 的形状广播到 [batch_size, len_q, len_k]，使其与查询序列的形状相匹配。这样可以在计算注意力分数时，按批次和序列长度进行操作。
    pad_attn_mask = ops.broadcast_to(pad_attn_mask, (batch_size, len_q, len_k))

    return pad_attn_mask


q = k = Tensor([[1, 3, 4, 5, 2, 0, 0]], mstype.float32)
pad_idx = 0  # 表示占位符<pad>的数字索引，即哪个数字代表<pad>
mask = get_attn_pad_mask(q, k, pad_idx)
print(mask)
print(q.shape, mask.shape)

attn = ops.ones(mask.shape)
print(attn)
attn = attn.masked_fill(mask, -1e9)
print(attn)
attn_weights = ops.softmax(attn)
print(attn_weights)


def get_attn_subsequent_mask(seq_q, seq_k):
    """生成时间掩码，使decoder在第t时刻只能看到序列的前t-1个元素

    Args:
        seq_q (Tensor): query序列，shape = [batch size, len_q]
        seq_k (Tensor): key序列，shape = [batch size, len_k]
    """
    batch_size, len_q = seq_q.shape
    batch_size, len_k = seq_k.shape
    # 生成一个全1的掩码，形状为 [batch_size, len_q, len_k]
    ones = ops.ones((batch_size, len_q, len_k), mstype.float32)
    # 将全1的掩码的下三角部分(含主对角线)置为0，形成一个上三角矩阵，用于屏蔽decoder在第t时刻只能看到序列的前t-1个元素
    subsequent_mask = mnp.triu(ones, k=1)
    return subsequent_mask


q = k = ops.ones((1, 5), mstype.float32)
mask = get_attn_subsequent_mask(q, k)
print(mask)